#include "sd_btree.h"
#include "sd_memory.h"
#include "sd_log.h"

static int    CreateNode(sd_BTree* tree, bool isLeaf);
static void   FreeNode(sd_BTreeNode* node);
static void   ZeroNode(sd_BTree* tree, int nodeIndex);
static int    TraverseNode(sd_BTree* tree, int nodeIndex, const sd_Key* key);
static sd_Ret SplitNode(sd_BTree* tree, int nodeIndex);
static sd_Ret RemoveNode(sd_BTree* tree, int nodeIndex);
static sd_Ret FindMedianKeyOfNode(sd_BTree* tree, int nodeIndex, sd_Key* medianKey);
static sd_Ret FindRecordOfNode(sd_BTree* tree, int nodeIndex, const sd_Key* key, sd_Record* record);
static sd_Ret DeleteRecordOfNode(sd_BTree* tree, int nodeIndex, const sd_Record* record);
static void   GrowNodeDepth(sd_BTree* tree, int nodeIndex);
static void   UpdateNodeDepth(sd_BTree* tree, int nodeIndex);
static void   PrintNode(sd_BTree* tree, int nodeIndex, FILE* fp, int indent);
static void   RebalanceTree(sd_BTree* tree);

sd_Ret sd_BTreeCreate(sd_BTree** tree, sd_KeyCompare key_cmp, int maxNumNodes, int maxChild, 
		int maxDepth, size_t recordsNodeBufSize)
{
	size_t size = sizeof(sd_BTree) + (maxNumNodes-1)*sizeof(sd_BTreeNode);
	*tree = (sd_BTree*)sd_malloc(size);
	if (*tree == NULL)
	{
		sd_log_err("Cannot allocate enough memory!\n");
		SD_ASSERT(false);
		return SD_RET_ALLOCATE_MEMEORY;
	}
	memset((*tree)->nodes, 0, maxNumNodes*sizeof(sd_BTreeNode));
	(*tree)->size = size;
	(*tree)->maxNumNodes = maxNumNodes;
	(*tree)->maxNumChildren = maxChild;
	(*tree)->recordsNodeBufSize = recordsNodeBufSize;
	(*tree)->maxDepth = maxDepth;
	(*tree)->rootNodeIndex = -1;
	(*tree)->depth = 0;
	(*tree)->key_cmp = key_cmp;
	return SD_RET_OK;
}

sd_Ret sd_BTreeFree(sd_BTree* tree)
{
	int i;
	SD_ASSERT(tree != NULL);
	for (i=0; i<tree->maxNumNodes; i++)
	{
		if (tree->nodes[i].isUsed)
		{
			FreeNode(&tree->nodes[i]);
		}
	}
	sd_free(tree);
	return SD_RET_OK;
}

sd_Ret sd_BTreeInsert(sd_BTree* tree, const sd_Record* record)
{
	sd_Ret ret;
	int nodeIndex;
	sd_BTreeNode* node;

	if (tree->rootNodeIndex == -1) /* No root node */
	{
		/* Create the root node */
		nodeIndex = CreateNode(tree, true);
		SD_ASSERT(nodeIndex != -1);
		node = &tree->nodes[nodeIndex];
		node->key = record->key;
		node->depth = 0; /* root node */
		node->isLeaf = true;
		sd_BlistInsert(&node->records, record);
		tree->rootNodeIndex = node->nodeIndex;
	}
	else
	{
		/* Find the node to insert the record */
		nodeIndex = TraverseNode(tree, tree->rootNodeIndex, &record->key);
		node = &tree->nodes[nodeIndex];

		/* Insert the record to the node's records list */
		ret = sd_BlistInsert(&node->records, record);
		SD_ON_FAILED_RETURN(ret);

		/* Check the condition for leaf node's splitting */
		if (node->isLeaf)
		{
			if (node->records.totalcount >= (size_t)tree->maxNumChildren)
			{
				ret = SplitNode(tree, node->nodeIndex);
				SD_ON_FAILED_RETURN(ret);
			}
		}
	}
	return SD_RET_OK;
}


sd_Ret sd_BTreeDelete(sd_BTree* tree, const sd_Record* record)
{
	sd_Ret ret;
	int nodeIndex;
	sd_BTreeNode* node;

	if (tree->rootNodeIndex == -1) /* No root node */
	{
		sd_log_err("This b-tree is empty!\n");
		SD_ASSERT(false);
		return SD_RET_DATA_EMPTY;
	}
	else
	{
		/* Find the node to insert the record */
		nodeIndex = TraverseNode(tree, tree->rootNodeIndex, &record->key);
		node = &tree->nodes[nodeIndex];

		/* Delete the record */
		ret = DeleteRecordOfNode(tree, nodeIndex, record);
		SD_ON_FAILED_RETURN(ret);

		/* Check this node if it is empty */
		if (node->records.totalcount == 0)
		{
			/* Remove this node */
			RemoveNode(tree, nodeIndex);
		}
	}
	return SD_RET_OK;
}

sd_Ret sd_BTreeSearch(sd_BTree* tree, const sd_Key* key, sd_Record* record)
{
	int nodeIndex;
	sd_BTreeNode* node;
	if (tree->rootNodeIndex == -1) /* No root node */
	{
		sd_log_err("This is no root node in the B-Tree indexing!\n");
		SD_ASSERT(false);
		return SD_RET_DATA_EMPTY;
	}

	/* Find the node to insert the record */
	nodeIndex = TraverseNode(tree, tree->rootNodeIndex, &record->key);
	node = &tree->nodes[nodeIndex];
	/* Get the record with the key */
	return FindRecordOfNode(tree, nodeIndex, key, record);
}


sd_Ret sd_BTreePrint(sd_BTree* tree, FILE* fp)
{
	PrintNode(tree, tree->rootNodeIndex, fp, 0);
	return SD_RET_OK;
}

static int CreateNode(sd_BTree* tree, bool isLeaf)
{
	sd_Ret ret;
	int newNodeIndex = -1;
	int i;
	sd_BTreeNode* node;
	for (i=0;i<tree->maxNumNodes; i++)
	{
		if (tree->nodes[i].isUsed == false)
		{
			newNodeIndex = i;
			break;
		}
	}
	if (i==tree->maxNumNodes)
	{
		/* No available node */
		return -1;
	}
	node = &tree->nodes[newNodeIndex];
	node->isUsed = true;
	node->depth = -1;
	node->isLeaf = isLeaf;
	node->nodeIndex = newNodeIndex;
	node->numChildNodes = 0;
	node->parentNodeIndex = -1;
	if (isLeaf)
	{
		ret = sd_BlistCreate(&node->records, sizeof(sd_Record),tree->recordsNodeBufSize,NULL);
		if (SD_FAILED(ret))
		{
			SD_ASSERT(false);
			return -1;
		}
	}

	return newNodeIndex;
}

static void FreeNode(sd_BTreeNode* node)
{
	sd_BlistFree(&node->records);
}

static void ZeroNode(sd_BTree* tree, int nodeIndex)
{
	sd_BTreeNode* node = &tree->nodes[nodeIndex];
	node->isUsed = false;
	node->depth = -1;
	node->nodeIndex = -1;
	node->parentNodeIndex = -1;
	node->numChildNodes = 0;
}

static int TraverseNode(sd_BTree* tree, int nodeIndex, const sd_Key* key)
{
	sd_BTreeNode* node;
	int childNodeIndex;
	sd_BTreeNode* child;
	int lb,ub;
	int cmp_ret; 

	node = &tree->nodes[nodeIndex];
	/* If it is a leaf node */
	if (node->isLeaf)
	{
		return node->nodeIndex;
	}

	cmp_ret = tree->key_cmp(key, &node->key);
	/* Binary search child nodes */
	
	lb = 0;
	ub = node->numChildNodes-1;
	/* Compare the current with the largest node's key */
	child = &tree->nodes[ub];
	cmp_ret = tree->key_cmp(key, &child->key);
	if (cmp_ret >= 0) /* Larger than largest node's key */
	{
		lb = ub;
	}
	else /* Less than the largest node's key */
	{
		while (lb+1 < ub)
		{
			childNodeIndex = node->childNodeIndexes[(lb+ub)/2];
			child = &tree->nodes[childNodeIndex];
			cmp_ret = tree->key_cmp(key, &child->key);
			if (cmp_ret >= 0) /* Larger or equal */
			{
				lb = (lb+ub)/2; 
			}
			else /* Less */
			{
				ub = (lb+ub)/2;
			}
		}
	}
	return  TraverseNode(tree, node->childNodeIndexes[lb], key);

}

static sd_Ret SplitNode(sd_BTree* tree, int nodeIndex)
{
	sd_Ret ret;
	sd_BTreeNode* node = &tree->nodes[nodeIndex];
	sd_BTreeNode* newNode;
	sd_BTreeNode* parentNode;
	int newNodeIndex;
	int parentNodeIndex = node->parentNodeIndex;
	int i;

	if (parentNodeIndex == -1 ) /* If it is the root node, which has no parent node */
	{
		parentNodeIndex = CreateNode(tree, false);
		SD_ASSERT(parentNodeIndex != -1);
		parentNode = &tree->nodes[parentNodeIndex];
		parentNode->isLeaf = false;
		parentNode->childNodeIndexes[parentNode->numChildNodes++] = nodeIndex;
		parentNode->key = node->key;
		parentNode->depth = 0; /* root node */
		parentNode->parentNodeIndex = -1; /* It shoule be the root node */
		node->parentNodeIndex = parentNodeIndex;
		tree->rootNodeIndex = parentNode->nodeIndex;
		/* Increase all the other nodes' depth */
		GrowNodeDepth(tree, nodeIndex);
	}

	/* Create a new node */
	newNodeIndex = CreateNode(tree, node->isLeaf);
	SD_ASSERT(newNodeIndex != -1);
	newNode = &tree->nodes[newNodeIndex];
	newNode->isLeaf = node->isLeaf;
	newNode->parentNodeIndex = node->parentNodeIndex;
	newNode->depth = node->depth;

	/* Get the common parent node */
	parentNode = &tree->nodes[node->parentNodeIndex];

	if (node->isLeaf)
	{
		sd_Key medianKey;
		sd_Record record;
		sd_BlistIterator iter;
		sd_Blist newRecordsOfOriginalNode;
		int cmp_ret;

		/* Find the median key of the records */
		ret = FindMedianKeyOfNode(tree, nodeIndex, &medianKey);
		SD_ON_FAILED_RETURN(ret);
		
		/* Create new record list for original node */
		ret = sd_BlistCreate(&newRecordsOfOriginalNode, sizeof(sd_Record),node->records.node_cap, NULL);
		SD_ON_FAILED_RETURN(ret);

		/* Split the record list of the node */
		ret = sd_BlistBegin(&node->records, &iter);
		SD_ON_FAILED_RETURN(ret);
		while (sd_BlistNext(&iter) != SD_RET_FALSE)
		{
			sd_BlistGet(&iter, &record);
			cmp_ret = tree->key_cmp(&medianKey, &record.key);
			if (cmp_ret <= 0) /* current key is larger than median key  */
			{
				sd_BlistInsert(&newNode->records, &record);
				if (cmp_ret == 0)
					newNode->key = record.key;
			}
			else
			{
				sd_BlistInsert(&newRecordsOfOriginalNode, &record);
			}
			cmp_ret = tree->key_cmp(&node->key, &record.key);
			if (cmp_ret > 0) /* The original node's key is larger than the record's key */
			{
				node->key = record.key;
			}
		}
		
		/* Update record list of the original node */
		sd_BlistFree(&node->records);
		node->records = newRecordsOfOriginalNode;
	}
	else
	{
		sd_BTreeNode* childNode;
		int medianindex = node->numChildNodes/2;
		int childindex;
		for (childindex=medianindex; childindex < node->numChildNodes; childindex++)
		{
			childNode = &tree->nodes[node->childNodeIndexes[childindex]];
			newNode->childNodeIndexes[newNode->numChildNodes++] = childNode->nodeIndex;
			childNode->parentNodeIndex = newNode->nodeIndex;
		}
		node->numChildNodes = (size_t)medianindex;
		childNode = &tree->nodes[node->childNodeIndexes[medianindex]];
		newNode->key = childNode->key;
	}

	/* Find the correct place and insert this node to parent node */
	for (i=0; i<parentNode->numChildNodes; i++)
	{
		if (parentNode->childNodeIndexes[i] == node->nodeIndex) /* Add the new node to the original node's right */
		{
			size_t movecount;
			movecount = parentNode->numChildNodes - i -1;
			if (movecount > 0)
				memmove(&parentNode->childNodeIndexes[i+2],&parentNode->childNodeIndexes[i+1], movecount*sizeof(int));
			parentNode->childNodeIndexes[i+1] = newNode->nodeIndex;
			parentNode->numChildNodes++;
			break;
		}
	}
	newNode->parentNodeIndex = parentNode->nodeIndex;
	/* Check the splitting condition of the parent node */
	if (parentNode->numChildNodes == tree->maxNumChildren)
	{
		ret = SplitNode(tree, parentNode->nodeIndex);
		SD_ON_FAILED_RETURN(ret);
	}

	return SD_RET_OK;
}

static sd_Ret RemoveNode(sd_BTree* tree, int nodeIndex)
{
	int i;
	sd_Ret ret;
	sd_BTreeNode* node = &tree->nodes[nodeIndex];
	sd_BTreeNode* parentNode = NULL;
	int depth = node->depth;

	if (node->isLeaf)
	{
		/* Free the record list */
		ret = sd_BlistFree(&node->records);
		SD_ON_FAILED_RETURN(ret);
	}

	/* Remove this node from its parent node */
	if (node->parentNodeIndex != -1)
	{
		parentNode = &tree->nodes[node->parentNodeIndex];
		for (i=0; i<parentNode->numChildNodes; i++)
		{
			if (parentNode->childNodeIndexes[i] == nodeIndex)
			{
				size_t movecount = (size_t)(parentNode->numChildNodes - i - 1);
				if (movecount > 0)
					memmove(&parentNode->childNodeIndexes[i], &parentNode->childNodeIndexes[i+1], movecount*sizeof(int));
				parentNode->numChildNodes--;
				break;
			}
		}
	}
	else
	{
		/* This is the root node to be removed */
		tree->rootNodeIndex = -1;
	}

	/* Zero this node */
	ZeroNode(tree, nodeIndex);

	/* Check its parent node need to remove again recursively */
	if (parentNode != NULL && parentNode->numChildNodes == 0)
	{
		ret = RemoveNode(tree, parentNode->nodeIndex);
		SD_ON_FAILED_RETURN(ret);
	}

	/* Update the tree's depth */
	if (depth == tree->depth)
	{
		if (tree->rootNodeIndex != -1)
			UpdateNodeDepth(tree, tree->rootNodeIndex);
		else
			tree->depth = -1;
	}

	return SD_RET_OK;
}

static sd_Ret FindMedianKeyOfNode(sd_BTree* tree, int nodeIndex, sd_Key* medianKey)
{
	sd_BTreeNode* node = &tree->nodes[nodeIndex];
	if (node->isLeaf) /* It is a leaf node */
	{
		/* Find the median key of the records */
		sd_Key key1, key2;
		sd_BlistIterator iter1, iter2;
		int count_less,count_equal,count_larger;
		int i;
		int cmp_ret;
		int totalcount = (int)(node->records.totalcount);
		int halfcount;
		sd_BlistBegin(&node->records, &iter1);
		for (i=0; i<totalcount; i++)
		{
			sd_BlistNext(&iter1);
			sd_BlistPartialGet(&iter1, &key1, 0, sizeof(sd_Key));
			count_less = 0;
			count_equal = 0;
			count_larger = 0;
			sd_BlistBegin(&node->records, &iter2);
			while (sd_BlistNext(&iter2) != SD_RET_FALSE)
			{
				sd_BlistPartialGet(&iter2, &key2, 0, sizeof(sd_Key));
				cmp_ret = tree->key_cmp(&key1, &key2);
				if (cmp_ret <0) 
					count_less++;
				else if (cmp_ret >0)
					count_larger++;
				else
					count_equal++;
			}

			/* Check it is a median key */
			halfcount = (totalcount - count_equal)/2;
			if (halfcount == count_less || halfcount == count_larger)
			{
				/* Found the median key */
				*medianKey = key1;
				break;
			}
		}
	}
	else /* Not a leaf node */
	{
		sd_BTreeNode* medianNode;
		int medianNodeIndex;
		medianNodeIndex = node->childNodeIndexes[node->numChildNodes/2];
		medianNode = &tree->nodes[medianNodeIndex];
		*medianKey = medianNode->key;
	}
	return SD_RET_OK;
}

static sd_Ret FindRecordOfNode(sd_BTree* tree, int nodeIndex, const sd_Key* key, sd_Record* record)
{
	sd_BTreeNode* node = &tree->nodes[nodeIndex];
	sd_BlistIterator iter;
	sd_Key tKey;
	int cmp_ret;

	sd_BlistBegin(&node->records, &iter);
	while (sd_BlistNext(&iter) != SD_RET_FALSE)
	{
		sd_BlistPartialGet(&iter, &tKey, 0, sizeof(sd_Key));
		cmp_ret = tree->key_cmp(key, &tKey);
		if (cmp_ret == 0)
		{
			sd_BlistGet(&iter, record);
			return SD_RET_TRUE; /* found */
		}
	}
	return SD_RET_KEY_NOTFOUND; /* Not found */
}

static sd_Ret DeleteRecordOfNode(sd_BTree* tree, int nodeIndex, const sd_Record* record)
{
	sd_Ret ret;
	sd_BTreeNode* node = &tree->nodes[nodeIndex];
	sd_BlistIterator iter;
	sd_Key tKey;
	sd_Record tRecord;
	int cmp_ret;
	bool bSpecifiedPayload = strlen(record->payload) > 0;
	int numOfDeleted = 0;

	sd_BlistBegin(&node->records, &iter);
	ret = sd_BlistNext(&iter);
	while (ret != SD_RET_FALSE)
	{
		sd_BlistPartialGet(&iter, &tKey, 0, sizeof(sd_Key));
		cmp_ret = tree->key_cmp(&record->key, &tKey);
		if (cmp_ret == 0)
		{
			if (bSpecifiedPayload)
			{
				sd_BlistGet(&iter, &tRecord);
				if (strcmp(tRecord.payload, record->payload) == 0)
				{
					numOfDeleted++;
					sd_BlistRemove(&iter);
					if (sd_BlistHasMore(&iter) == SD_RET_FALSE)
						break;
					else
						continue; /* Not have to go next for this iterator */
				}
			}
			else
			{
				numOfDeleted++;
				sd_BlistRemove(&iter);
				if (sd_BlistHasMore(&iter) == SD_RET_FALSE)
					break;
				else
					continue; /* Not have to go next for this iterator */
			}
		}
		/* Jump to the next record in the list */
		ret = sd_BlistNext(&iter);
	}

	if (numOfDeleted > 0)
		return SD_RET_OK;
	else
	{
		/* Not found */
		if (bSpecifiedPayload)
			return SD_RET_RECORD_NOTFOUND;
		else
			return SD_RET_KEY_NOTFOUND;
	}	
}


static void GrowNodeDepth(sd_BTree* tree, int nodeIndex)
{
	int i;
	sd_BTreeNode* node = &tree->nodes[nodeIndex];
	node->depth++;
	if (node->isLeaf)
	{
		/* Keep the tree's depth be the maximum depth among nodes */
		if (node->depth > tree->depth)
			tree->depth = node->depth;
	}
	else
	{
		for (i=0;i<node->numChildNodes;i++)
		{
			GrowNodeDepth(tree, node->childNodeIndexes[i]);
		}
	}
}

static void UpdateNodeDepth(sd_BTree* tree, int nodeIndex)
{
	int i;
	sd_BTreeNode* node = &tree->nodes[nodeIndex];
	if (node->isLeaf)
	{
		/* Keep the tree's depth be the maximum depth among nodes */
		if (node->depth > tree->depth)
			tree->depth = node->depth;
	}
	else
	{
		for (i=0;i<node->numChildNodes;i++)
		{
			UpdateNodeDepth(tree, node->childNodeIndexes[i]);
		}
	}
}

static void RebalanceTree(sd_BTree* tree)
{
	sd_BTreeNode* rootNode = &tree->nodes[tree->rootNodeIndex];

}

static void PrintNode(sd_BTree* tree, int nodeIndex, FILE* fp, int indent)
{
	sd_BTreeNode* node = &tree->nodes[nodeIndex];
	int i;
	sd_BlistIterator iter;
	sd_Key key;

	for (i=0; i<indent; i++)
		fprintf(fp, "\t");
	fprintf(fp,"Node (key=");
	sd_KeyPrint(&node->key, fp);
	fprintf(fp,",idx=%d,pi=%d), count=%d : ", node->nodeIndex,node->parentNodeIndex,node->records.totalcount);
	sd_BlistBegin(&node->records, &iter);
	for (i=0; i<(int)node->records.totalcount; i++)
	{
		sd_BlistNext(&iter);
		sd_BlistPartialGet(&iter,&key, 0, sizeof(sd_Key)); 
		sd_KeyPrint(&key, fp);
		if (i < (int)node->records.totalcount-1)
			fprintf(fp,", ");
	}
	fprintf(fp,"\n");
	for (i=0; i<node->numChildNodes; i++)
	{
		PrintNode(tree, node->childNodeIndexes[i], fp, indent+1);
	}
}

